#include <api.h>

typedef struct Trait {
    float value;
} Trait;

typedef struct TraitA {
    float value;
} TraitA;

typedef struct TraitB {
    float value;
} TraitB;

void ProcessTraits(ecs_iter_t *it) {
    Trait *tr = ecs_column(it, Trait, 1);

    probe_system(it);

    if (tr) {
        int32_t i;
        for (i = 0; i < it->count; i ++) {
            tr[i].value ++;
        }
    }
}

void Traits_type_w_one_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_SYSTEM(world, ProcessTraits, EcsOnUpdate, TRAIT | Trait);

    /* Ensure that trait is matched against different components */
    ecs_entity_t e1 = ecs_set_trait(world, 0, Position, Trait, {
        .value = 10
    });  
    test_assert(e1 != 0);
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(Trait)));

    ecs_entity_t e2 = ecs_set_trait(world, 0, Velocity, Trait, {
        .value = 20
    });
    test_assert(e2 != 0);
    test_assert( ecs_has_trait(world, e2, ecs_typeid(Velocity), ecs_typeid(Trait)));

    Probe ctx = {0};
    ecs_set_context(world, &ctx);

    ecs_progress(world, 0);

    const Trait* tr = ecs_get_trait(world, e1, Position, Trait);
    test_assert(tr != NULL);
    test_int(tr->value, 11);

    tr = ecs_get_trait(world, e2, Velocity, Trait);
    test_assert(tr != NULL);
    test_int(tr->value, 21); 

    test_int(ctx.count, 2);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, ProcessTraits);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);

    test_int(ctx.e[0], e1);
    test_int(ctx.e[1], e2);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Position));

    c = ctx.c[1][0] & ECS_COMPONENT_MASK;
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Velocity));

    test_int(ctx.s[0][0], 0);
    test_int(ctx.s[1][0], 0);

    ecs_fini(world);
}

void Traits_type_w_two_traits() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_SYSTEM(world, ProcessTraits, EcsOnUpdate, TRAIT | Trait);

    /* Ensure that trait is matched against different components on same entity */
    ecs_entity_t e1 = ecs_new(world, 0);
    ecs_set_trait(world, e1, Position, Trait, {
        .value = 10
    });  
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(Trait)));

    ecs_set_trait(world, e1, Velocity, Trait, {
        .value = 20
    });      
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Velocity), ecs_typeid(Trait)));

    ecs_entity_t e2 = ecs_new(world, 0);
    ecs_set_trait(world, e2, Position, Trait, {
        .value = 30
    });
    test_assert( ecs_has_trait(world, e2, ecs_typeid(Position), ecs_typeid(Trait)));

    ecs_set_trait(world, e2, Velocity, Trait, {
        .value = 40
    });
    test_assert( ecs_has_trait(world, e2, ecs_typeid(Position), ecs_typeid(Trait)));

    Probe ctx = {0};
    ecs_set_context(world, &ctx);

    ecs_progress(world, 0);

    const Trait* tr = ecs_get_trait(world, e1, Position, Trait);
    test_assert(tr != NULL);
    test_int(tr->value, 11);

    tr = ecs_get_trait(world, e1, Velocity, Trait);
    test_assert(tr != NULL);
    test_int(tr->value, 21);

    tr = ecs_get_trait(world, e2, Position, Trait);
    test_assert(tr != NULL);
    test_int(tr->value, 31);

    tr = ecs_get_trait(world, e2, Velocity, Trait);
    test_assert(tr != NULL);
    test_int(tr->value, 41); 

    test_int(ctx.count, 4);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, ProcessTraits);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);

    test_int(ctx.e[0], e1);
    test_int(ctx.e[1], e2);
    test_int(ctx.e[2], e1);
    test_int(ctx.e[3], e2);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Velocity));

    c = ctx.c[1][0] & ECS_COMPONENT_MASK;
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Position));

    test_int(ctx.s[0][0], 0);
    test_int(ctx.s[1][0], 0);
    test_int(ctx.s[2][0], 0);
    test_int(ctx.s[3][0], 0);

    ecs_fini(world);
}

void Traits_add_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_SYSTEM(world, ProcessTraits, EcsOnUpdate, TRAIT | Trait);

    ecs_entity_t e1 = ecs_new(world, 0);
    test_assert(e1 != 0);

    ecs_add_trait(world, e1, ecs_typeid(Position), ecs_typeid(Trait));
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(Trait)));

    ecs_fini(world);
}

void Traits_remove_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_SYSTEM(world, ProcessTraits, EcsOnUpdate, TRAIT | Trait);

    ecs_entity_t e1 = ecs_set_trait(world, 0, Position, Trait, {
        .value = 10
    });  
    test_assert(e1 != 0);
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(Trait)));

    ecs_remove_trait(world, e1, ecs_typeid(Position), ecs_typeid(Trait));

    test_assert( !ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(Trait)));

    ecs_fini(world);
}

void Traits_add_tag_trait_for_tag() {
    ecs_world_t *world = ecs_init();

    ECS_TAG(world, Tag1);
    ECS_TAG(world, Tag2);
    ECS_TAG(world, TraitTag);

    ECS_SYSTEM(world, ProcessTraits, EcsOnUpdate, TRAIT | TraitTag);

    /* Ensure that trait is matched against different components */
    ecs_entity_t e1 = ecs_new(world, 0);
    ecs_add_entity(world, e1, ecs_trait(Tag1, TraitTag));
    test_assert(e1 != 0);
    test_assert( ecs_has_entity(world, e1, ecs_trait(Tag1, TraitTag)));

    ecs_entity_t e2 = ecs_new(world, 0);
    ecs_add_entity(world, e2, ecs_trait(Tag2, TraitTag));
    test_assert(e2 != 0);
    test_assert( ecs_has_entity(world, e2, ecs_trait(Tag2, TraitTag)));

    Probe ctx = {0};
    ecs_set_context(world, &ctx);

    ecs_progress(world, 0);

    test_int(ctx.count, 2);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, ProcessTraits);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);

    test_int(ctx.e[0], e1);
    test_int(ctx.e[1], e2);

    ecs_entity_t c = ctx.c[0][0];

    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, TraitTag);
    test_int(lo, Tag1);

    c = ctx.c[1][0] & ECS_COMPONENT_MASK;
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, TraitTag);
    test_int(lo, Tag2);

    test_int(ctx.s[0][0], 0);
    test_int(ctx.s[1][0], 0);

    ecs_fini(world);
}

void ProcessValueTraits(ecs_iter_t *it) {
    /* Strictly speaking this can be either Position or Velocity, but they have
     * the same layout. */
    Position *p = ecs_column(it, Position, 1);

    probe_system(it);

    if (p) {
        int32_t i;
        for (i = 0; i < it->count; i ++) {
            p[i].x += 10;
            p[i].y += 20;
        }
    }
}

void Traits_add_tag_trait_for_component() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_TAG(world, TraitTag);

    ECS_SYSTEM(world, ProcessValueTraits, EcsOnUpdate, TRAIT | TraitTag);

    ecs_entity_t e1 = ecs_set_trait_tag(world, 0, TraitTag, Position, {
        .x = 1,
        .y = 2
    });
    test_assert( ecs_has_entity(world, e1, ecs_trait(ecs_typeid(Position), TraitTag)));

    ecs_entity_t e2 = ecs_set_trait_tag(world, 0, TraitTag, Velocity, {
        .x = 3,
        .y = 4
    });
    test_assert( ecs_has_entity(world, e2, ecs_trait(ecs_typeid(Velocity), TraitTag)));

    Probe ctx = {0};
    ecs_set_context(world, &ctx);

    ecs_progress(world, 0);

    const Position* tr_p = ecs_get_w_entity(world, e1, ecs_trait(ecs_typeid(Position), TraitTag));
    test_assert(tr_p != NULL);
    test_int(tr_p->x, 11);
    test_int(tr_p->y, 22);

    const Velocity* tr_v = ecs_get_w_entity(world, e2, ecs_trait(ecs_typeid(Velocity), TraitTag));
    test_assert(tr_v != NULL);
    test_int(tr_v->x, 13);
    test_int(tr_v->y, 24);

    test_int(ctx.count, 2);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, ProcessValueTraits);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);

    test_int(ctx.e[0], e1);
    test_int(ctx.e[1], e2);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, TraitTag);
    test_int(lo, ecs_typeid(Position));

    c = ctx.c[1][0] & ECS_COMPONENT_MASK;
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, TraitTag);
    test_int(lo, ecs_typeid(Velocity));

    test_int(ctx.s[0][0], 0);
    test_int(ctx.s[1][0], 0);

    ecs_fini(world);
}

void ProcessTwoTraits(ecs_iter_t *it) {
    TraitA *tr_a = ecs_column(it, TraitA, 1);
    TraitB *tr_b = ecs_column(it, TraitB, 2);

    probe_system(it);

    int32_t i;
    for (i = 0; i < it->count; i ++) {
        tr_a[i].value ++;
        tr_b[i].value ++;
    }
}

void Traits_query_2_traits() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, TraitA);
    ECS_COMPONENT(world, TraitB);

    ECS_SYSTEM(world, ProcessTwoTraits, EcsOnUpdate, TRAIT | TraitA, TRAIT | TraitB);

    /* Create entity with both TraitA and TraitB */
    ecs_entity_t e1 = ecs_set_trait(world, 0, Position, TraitA, {
        .value = 1
    });

    ecs_set_trait(world, e1, Position, TraitB, {
        .value = 2
    });    

    test_assert( ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(TraitA)));
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(TraitB)));

    /* Create entity with only TraitA. Should not be matched with system */
    ecs_entity_t e2 = ecs_set_trait(world, 0, Position, TraitA, {
        .value = 3
    });
    test_assert( ecs_has_trait(world, e2, ecs_typeid(Position), ecs_typeid(TraitA)));

    /* Run system */
    Probe ctx = {0};
    ecs_set_context(world, &ctx);
    ecs_progress(world, 0);

    const TraitA *tr_a = ecs_get_trait(world, e1, Position, TraitA);
    test_int(tr_a->value, 2);

    const TraitB *tr_b = ecs_get_trait(world, e1, Position, TraitB);
    test_int(tr_b->value, 3);

    tr_a = ecs_get_trait(world, e2, Position, TraitA);
    test_int(tr_a->value, 3);

    test_int(ctx.count, 1);
    test_int(ctx.invoked, 1);
    test_int(ctx.system, ProcessTwoTraits);
    test_int(ctx.column_count, 2);
    test_null(ctx.param);

    test_int(ctx.e[0], e1);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(TraitA));
    test_int(lo, ecs_typeid(Position));

    c = ctx.c[0][1];
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(TraitB));
    test_int(lo, ecs_typeid(Position));    

    ecs_fini(world);
}

void Traits_query_2_traits_2_instances_per_type() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, TraitA);
    ECS_COMPONENT(world, TraitB);

    ECS_SYSTEM(world, ProcessTwoTraits, EcsOnUpdate, TRAIT | TraitA, TRAIT | TraitB);

    /* Create entity with both TraitA and TraitB, applied to two components*/
    ecs_entity_t e1 = ecs_set_trait(world, 0, Position, TraitA, {
        .value = 1
    });

    ecs_set_trait(world, e1, Position, TraitB, {
        .value = 2
    });

    ecs_set_trait(world, e1, Velocity, TraitA, {
        .value = 3
    });

    ecs_set_trait(world, e1, Velocity, TraitB, {
        .value = 4
    });        

    test_assert( ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(TraitA)));
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Position), ecs_typeid(TraitB)));
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Velocity), ecs_typeid(TraitA)));
    test_assert( ecs_has_trait(world, e1, ecs_typeid(Velocity), ecs_typeid(TraitB)));

    /* Run system */
    Probe ctx = {0};
    ecs_set_context(world, &ctx);
    ecs_progress(world, 0);

    const TraitA *tr_a = ecs_get_trait(world, e1, Position, TraitA);
    test_int(tr_a->value, 2);

    tr_a = ecs_get_trait(world, e1, Velocity, TraitA);
    test_int(tr_a->value, 4);    

    const TraitB *tr_b = ecs_get_trait(world, e1, Position, TraitB);
    test_int(tr_b->value, 3);

    tr_b = ecs_get_trait(world, e1, Velocity, TraitB);
    test_int(tr_b->value, 5);

    test_int(ctx.count, 2);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, ProcessTwoTraits);
    test_int(ctx.column_count, 2);
    test_null(ctx.param);

    test_int(ctx.e[0], e1);
    test_int(ctx.e[1], e1);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(TraitA));
    test_int(lo, ecs_typeid(Position));

    c = ctx.c[0][1];
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(TraitB));
    test_int(lo, ecs_typeid(Position));    

    c = ctx.c[1][0];
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(TraitA));
    test_int(lo, ecs_typeid(Velocity));

    c = ctx.c[1][1];
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(TraitB));
    test_int(lo, ecs_typeid(Velocity));  

    ecs_fini(world);
}

void Traits_override_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Trait);

    ecs_entity_t base = ecs_new(world, 0);
    ecs_set_trait(world, base, Position, Trait, {.value = 10});

    ecs_entity_t instance = ecs_new_w_entity(world, ECS_INSTANCEOF | base);
    test_assert(ecs_has_trait(world, instance, ecs_typeid(Position), ecs_typeid(Trait)));

    const Trait *t = ecs_get_trait(world, instance, Position, Trait);
    test_assert(t != NULL);
    test_int(t->value, 10);

    const Trait *t_2 = ecs_get_trait(world, base, Position, Trait);
    test_assert(t_2 != NULL);
    test_assert(t == t_2);

    ecs_add_trait(world, instance, ecs_typeid(Position), ecs_typeid(Trait));
    t = ecs_get_trait(world, instance, Position, Trait);
    test_assert(t != NULL);
    test_int(t->value, 10);
    test_assert(t != t_2);

    ecs_fini(world);
}

void Traits_override_tag_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_TAG(world, Trait);

    ecs_entity_t base = ecs_new(world, 0);
    ecs_set_trait_tag(world, base, Trait, Position, {.x = 10, .y = 20});

    ecs_entity_t instance = ecs_new_w_entity(world, ECS_INSTANCEOF | base);
    test_assert(ecs_has_trait(world, instance, ecs_typeid(Position), Trait));

    const Position *t = ecs_get_trait_tag(world, instance, Trait, Position);
    test_assert(t != NULL);
    test_int(t->x, 10);
    test_int(t->y, 20);

    const Position *t_2 = ecs_get_trait_tag(world, base, Trait, Position);
    test_assert(t_2 != NULL);
    test_assert(t == t_2);

    ecs_add_trait(world, instance, ecs_typeid(Position), Trait);
    t = ecs_get_trait_tag(world, instance, Trait, Position);
    test_assert(t != NULL);
    test_int(t->x, 10);
    test_int(t->y, 20);
    test_assert(t != t_2);

    ecs_fini(world);
}

void Traits_trait_wildcard_system() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_SYSTEM(world, ProcessTraits, EcsOnUpdate, TRAIT | Trait, TRAIT | Trait > *);

    /* Ensure that trait is matched against different components */
    ecs_entity_t e1 = 
    ecs_set_trait(world, 0, Position, Trait, {
        .value = 10
    });
    ecs_add(world, e1, Position);

    ecs_entity_t e2 = 
    ecs_set_trait(world, 0, Velocity, Trait, {
        .value = 20
    });
    ecs_add(world, e2, Velocity);

    /* This entity should not be matched as it doesn't have Position */
    ecs_set_trait(world, 0, Position, Trait, {
        .value = 20
    });

    /* This entity should not be matched as it doesn't have the trait */
    ecs_set(world, 0, Position, { 10, 20 });

    Probe ctx = {0};
    ecs_set_context(world, &ctx);

    ecs_progress(world, 0);

    const Trait* tr = ecs_get_trait(world, e1, Position, Trait);
    test_assert(tr != NULL);
    test_int(tr->value, 11);

    tr = ecs_get_trait(world, e2, Velocity, Trait);
    test_assert(tr != NULL);
    test_int(tr->value, 21); 

    test_int(ctx.count, 2);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, ProcessTraits);
    test_int(ctx.column_count, 2);
    test_null(ctx.param);

    test_int(ctx.e[0], e1);
    test_int(ctx.e[1], e2);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Position));
    test_int(ctx.c[0][1], ecs_typeid(Position));

    c = ctx.c[1][0] & ECS_COMPONENT_MASK;
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Velocity));
    test_int(ctx.c[1][1], ecs_typeid(Velocity));

    test_int(ctx.s[0][0], 0);
    test_int(ctx.s[1][0], 0);

    ecs_fini(world);
}

static
void IterTraitWildcard(ecs_iter_t *it) {
    probe_system(it);
}

void Traits_trait_only_wildcard_system() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_SYSTEM(world, IterTraitWildcard, EcsOnUpdate, TRAIT | Trait > *);

    /* Ensure that trait is matched against different components */
    ecs_entity_t e1 = 
    ecs_set_trait(world, 0, Position, Trait, {
        .value = 10
    });
    ecs_add(world, e1, Position);

    ecs_entity_t e2 = 
    ecs_set_trait(world, 0, Velocity, Trait, {
        .value = 20
    });
    ecs_add(world, e2, Velocity);

    /* This entity should not be matched as it doesn't have Position */
    ecs_set_trait(world, 0, Position, Trait, {
        .value = 20
    });

    /* This entity should not be matched as it doesn't have the trait */
    ecs_set(world, 0, Position, { 10, 20 });

    Probe ctx = {0};
    ecs_set_context(world, &ctx);

    ecs_progress(world, 0);

    test_int(ctx.count, 2);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, IterTraitWildcard);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);

    test_int(ctx.e[0], e1);
    test_int(ctx.e[1], e2);

    test_int(ctx.c[0][0], ecs_typeid(Position));
    test_int(ctx.c[1][0], ecs_typeid(Velocity));

    test_int(ctx.s[0][0], 0);
    test_int(ctx.s[1][0], 0);

    ecs_fini(world);
}

static int set_trait_invoked = 0;

static
void SetTrait(ecs_iter_t *it) {
    set_trait_invoked ++;
    probe_system(it);
}

void Traits_trait_wildcard_on_set() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_SYSTEM(world, SetTrait, EcsOnSet, TRAIT | Trait > *);

    /* Ensure that trait is matched against different components */
    ecs_entity_t e1 = 
    ecs_set_trait(world, 0, Position, Trait, { .value = 10 });
    ecs_add(world, e1, Position);
    test_int(set_trait_invoked, 0);

    ecs_entity_t e2 = 
    ecs_set_trait(world, 0, Velocity, Trait, { .value = 20 });
    ecs_add(world, e2, Velocity);
    test_int(set_trait_invoked, 0);

    /* This entity should not trigger as it doesn't have the trait */
    ecs_set(world, 0, Position, { 10, 20 });
    test_int(set_trait_invoked, 0);

    /* Set Position on e1, should trigger OnSet */
    ecs_set(world, e1, Position, {10, 20});
    test_int(set_trait_invoked, 1);

    ecs_set(world, e2, Position, {10, 20});
    test_int(set_trait_invoked, 1);

    ecs_set(world, e2, Velocity, {10, 20});
    test_int(set_trait_invoked, 2);

    ecs_fini(world);
}

static
void TraitTrigger(ecs_iter_t *it) {
    probe_system(it);
}

void Traits_on_add_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_TRIGGER(world, TraitTrigger, EcsOnAdd, Trait);

    Probe ctx = { 0 };
    ecs_set_context(world, &ctx);

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_trait(world, e, ecs_typeid(Position), ecs_typeid(Trait));

    test_int(ctx.count, 1);
    test_int(ctx.invoked, 1);
    test_int(ctx.system, TraitTrigger);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);
    test_int(ctx.e[0], e);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Position));  

    ctx = (Probe){ 0 };
    ecs_add_trait(world, e, ecs_typeid(Velocity), ecs_typeid(Trait));

    test_int(ctx.count, 1);
    test_int(ctx.invoked, 1);
    test_int(ctx.system, TraitTrigger);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);
    test_int(ctx.e[0], e);

    c = ctx.c[0][0];
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Velocity));

    ecs_fini(world);
}

void Traits_on_add_trait_tag() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_TAG(world, Trait);

    ECS_TRIGGER(world, TraitTrigger, EcsOnAdd, Position);

    Probe ctx = { 0 };
    ecs_set_context(world, &ctx);

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_trait(world, e, ecs_typeid(Position), Trait);

    test_int(ctx.count, 1);
    test_int(ctx.invoked, 1);
    test_int(ctx.system, TraitTrigger);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);
    test_int(ctx.e[0], e);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, Trait);
    test_int(lo, ecs_typeid(Position));  

    ecs_fini(world);
}

void Traits_on_remove_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_TRIGGER(world, TraitTrigger, EcsOnRemove, Trait);

    Probe ctx = { 0 };
    ecs_set_context(world, &ctx);

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_trait(world, e, ecs_typeid(Position), ecs_typeid(Trait));
    ecs_add_trait(world, e, ecs_typeid(Velocity), ecs_typeid(Trait));
    test_int(ctx.count, 0);

    ecs_remove_trait(world, e, ecs_typeid(Position), ecs_typeid(Trait));
    test_int(ctx.count, 1);
    test_int(ctx.invoked, 1);
    test_int(ctx.system, TraitTrigger);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);
    test_int(ctx.e[0], e);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Position));  

    ctx = (Probe){ 0 };
    ecs_remove_trait(world, e, ecs_typeid(Velocity), ecs_typeid(Trait));
    test_int(ctx.count, 1);
    test_int(ctx.invoked, 1);
    test_int(ctx.system, TraitTrigger);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);
    test_int(ctx.e[0], e);

    c = ctx.c[0][0];
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Velocity));

    ecs_fini(world);
}

void Traits_on_remove_trait_tag() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_TAG(world, Trait);

    ECS_TRIGGER(world, TraitTrigger, EcsOnRemove, Position);

    Probe ctx = { 0 };
    ecs_set_context(world, &ctx);

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_trait(world, e, ecs_typeid(Position), Trait);
    test_int(ctx.count, 0);

    ecs_remove_trait(world, e, ecs_typeid(Position), Trait);    
    test_int(ctx.count, 1);
    test_int(ctx.invoked, 1);
    test_int(ctx.system, TraitTrigger);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);
    test_int(ctx.e[0], e);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, Trait);
    test_int(lo, ecs_typeid(Position));  

    ecs_fini(world);
}

void Traits_on_remove_trait_on_delete() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_COMPONENT(world, Trait);

    ECS_TRIGGER(world, TraitTrigger, EcsOnRemove, Trait);

    Probe ctx = { 0 };
    ecs_set_context(world, &ctx);

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_trait(world, e, ecs_typeid(Position), ecs_typeid(Trait));
    ecs_add_trait(world, e, ecs_typeid(Velocity), ecs_typeid(Trait));
    test_int(ctx.count, 0);

    ecs_delete(world, e);
    test_int(ctx.count, 2);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, TraitTrigger);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);
    test_int(ctx.e[0], e);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Position));  

    test_int(ctx.e[1], e);

    c = ctx.c[1][0];
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, ecs_typeid(Trait));
    test_int(lo, ecs_typeid(Velocity));

    ecs_fini(world);
}

static
void TraitTriggerPosition(ecs_iter_t *it) {
    TraitTrigger(it);
}

static
void TraitTriggerVelocity(ecs_iter_t *it) {
    TraitTrigger(it);
}

void Traits_on_remove_trait_tag_on_delete() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Velocity);
    ECS_TAG(world, Trait);

    ECS_TRIGGER(world, TraitTriggerPosition, EcsOnRemove, Position);
    ECS_TRIGGER(world, TraitTriggerVelocity, EcsOnRemove, Velocity);

    Probe ctx = { 0 };
    ecs_set_context(world, &ctx);

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_trait(world, e, ecs_typeid(Position), Trait);
    ecs_add_trait(world, e, ecs_typeid(Velocity), Trait);
    test_int(ctx.count, 0);

    ecs_delete(world, e);
    test_int(ctx.count, 2);
    test_int(ctx.invoked, 2);
    test_int(ctx.system, TraitTriggerVelocity);
    test_int(ctx.column_count, 1);
    test_null(ctx.param);
    test_int(ctx.e[0], e);

    ecs_entity_t c = ctx.c[0][0];
    ecs_entity_t hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    ecs_entity_t lo = ecs_entity_t_lo(c);
    test_int(hi, Trait);
    test_int(lo, ecs_typeid(Position));  

    test_int(ctx.e[1], e);

    c = ctx.c[1][0];
    hi = ecs_entity_t_hi(c & ECS_COMPONENT_MASK);
    lo = ecs_entity_t_lo(c);
    test_int(hi, Trait);
    test_int(lo, ecs_typeid(Velocity));

    ecs_fini(world);
}

void Traits_trait_from_shared() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Trait);

    ecs_entity_t base = ecs_new(world, 0);
    ecs_add_trait(world, base, ecs_typeid(Position), ecs_typeid(Trait));

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_entity(world, e, ECS_INSTANCEOF | base);

    ecs_query_t *q = ecs_query_new(world, "SHARED:TRAIT|Trait>Position"); // ew
    
    int32_t count = 0;
    ecs_iter_t it = ecs_query_iter(q);
    while (ecs_query_next(&it)) {
        Trait *t = ecs_column(&it, Trait, 1);
        test_assert(t != NULL);

        int i;
        for (i = 0; i < it.count; i ++) {
            test_assert(it.entities[i] == e);
        }
        count += it.count;
    }

    test_assert(count == 1);

    ecs_fini(world);
}

void Traits_simple_trait_syntax() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Trait);

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_trait(world, e, ecs_typeid(Position), ecs_typeid(Trait));

    ecs_query_t *q = ecs_query_new(world, "Trait FOR Position"); // not ew
    
    int32_t count = 0;
    ecs_iter_t it = ecs_query_iter(q);
    while (ecs_query_next(&it)) {
        int i;
        for (i = 0; i < it.count; i ++) {
            test_assert(it.entities[i] == e);
        }
        count += it.count;
    }

    test_assert(count == 1);

    ecs_fini(world);
}

void Traits_trait_w_component_query() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_COMPONENT(world, Trait);

    ecs_entity_t e = ecs_new(world, 0);
    ecs_add_trait(world, e, ecs_typeid(Position), ecs_typeid(Trait));

    ecs_query_t *q = ecs_query_new(world, "TRAIT | Trait > Position");
    
    int32_t count = 0;
    ecs_iter_t it = ecs_query_iter(q);
    while (ecs_query_next(&it)) {
        Trait *t = ecs_column(&it, Trait, 1);
        test_assert(t != NULL);

        int i;
        for (i = 0; i < it.count; i ++) {
            test_assert(it.entities[i] == e);
        }
        count += it.count;
    }

    test_assert(count == 1);

    ecs_fini(world);
}

void Traits_query_trait_or_component() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_TAG(world, Trait);

    ecs_entity_t e1 = ecs_new(world, 0);
    ecs_add_trait(world, e1, ecs_typeid(Position), Trait);
    ecs_entity_t e2 = ecs_new(world, Position);

    ecs_query_t *q = ecs_query_new(world, "TRAIT | Trait > Position || Position");
    
    int32_t count = 0;
    ecs_iter_t it = ecs_query_iter(q);
    while (ecs_query_next(&it)) {
        Position *t = ecs_column(&it, Position, 1);
        test_assert(t != NULL);

        int i;
        for (i = 0; i < it.count; i ++) {
            test_assert(it.entities[i] == e1 || it.entities[i] == e2);
        }
        count += it.count;
    }

    test_assert(count == 2);

    ecs_fini(world);
}

void Traits_query_trait_or_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_TAG(world, TraitA);
    ECS_TAG(world, TraitB);

    ecs_entity_t e1 = ecs_new(world, 0);
    ecs_add_trait(world, e1, ecs_typeid(Position), TraitA);
    ecs_entity_t e2 = ecs_new(world, Position);
    ecs_add_trait(world, e2, ecs_typeid(Position), TraitB);

    ecs_query_t *q = ecs_query_new(world, "TRAIT | TraitA > Position || TRAIT | TraitB > Position");
    
    int32_t count = 0;
    ecs_iter_t it = ecs_query_iter(q);
    while (ecs_query_next(&it)) {
        Position *t = ecs_column(&it, Position, 1);
        test_assert(t != NULL);

        int i;
        for (i = 0; i < it.count; i ++) {
            test_assert(it.entities[i] == e1 || it.entities[i] == e2);
        }
        count += it.count;
    }

    test_assert(count == 2);

    ecs_fini(world);
}

void Traits_query_not_trait() {
    ecs_world_t *world = ecs_init();

    ECS_COMPONENT(world, Position);
    ECS_TAG(world, Trait);

    ecs_entity_t e1 = ecs_new(world, Position);
    ecs_entity_t e2 = ecs_new(world, Position);
    ecs_add_trait(world, e2, ecs_typeid(Position), Trait);

    ecs_query_t *q = ecs_query_new(world, "!TRAIT | Trait > Position, Position");
    
    int32_t count = 0;
    ecs_iter_t it = ecs_query_iter(q);
    while (ecs_query_next(&it)) {
        Position *t = ecs_column(&it, Position, 1);
        test_assert(t == NULL);
        Position *p = ecs_column(&it, Position, 2);
        test_assert(p != NULL);

        int i;
        for (i = 0; i < it.count; i ++) {
            test_assert(it.entities[i] == e1);
        }
        count += it.count;
    }

    test_assert(count == 1);

    ecs_fini(world);
}
